import graphviz
import os
from typing import Dict, Set, Tuple

# 设置Graphviz路径
GRAPHVIZ_PATH = r'E:\Graphviz-14.0.4-win64\bin'
os.environ["PATH"] += os.pathsep + GRAPHVIZ_PATH

class NFAVisualizer:
    def __init__(self):
        self.dot = graphviz.Digraph('NFA', comment='Non-deterministic Finite Automaton')
        self.dot.attr(rankdir='LR', size='12,10', splines='true')
        
    @staticmethod
    def parse_states(states_str: str) -> Set[int]:
        """解析状态集合字符串"""
        clean_str = states_str.strip('{} ')
        if not clean_str:
            return set()
        return {int(s.strip()) for s in clean_str.split(',') if s.strip()}
    
    @staticmethod
    def parse_alphabet(alphabet_str: str) -> Set[str]:
        """解析字母表字符串"""
        clean_str = alphabet_str.strip('{} ')
        if not clean_str:
            return set()
        return {s.strip() for s in clean_str.split(',') if s.strip()}
    
    @staticmethod
    def parse_start_state(start_state_str: str) -> int:
        """解析起始状态字符串"""
        clean_str = start_state_str.strip('{} ')
        return int(clean_str)
    
    @staticmethod
    def parse_accept_states(accept_states_str: str) -> Set[int]:
        """解析接受状态集合字符串"""
        return NFAVisualizer.parse_states(accept_states_str)
    
    def visualize_from_data(self, 
                           states: Set[int],
                           alphabet: Set[str],
                           transitions: Dict[Tuple[int, str], Set[int]],
                           start_state: int,
                           accept_states: Set[int],
                           output_file: str = 'nfa'):
        """从数据直接可视化NFA"""
        self._create_graph(states, alphabet, transitions, start_state, accept_states)
        return self._render(output_file)
    
    def _create_graph(self, states, alphabet, transitions, start_state, accept_states):
        """创建NFA图形"""
        self.dot = graphviz.Digraph('NFA', comment='Non-deterministic Finite Automaton')
        self.dot.attr(rankdir='LR', size='12,10', splines='true', nodesep='0.5', ranksep='1')
        
        # 添加起始箭头
        self.dot.node('start', shape='none', label='', width='0.1', height='0.1')
        self.dot.edge('start', f'{start_state}', arrowhead='normal')
        
        # 添加状态节点
        for state in states:
            if state in accept_states:
                self.dot.node(str(state), shape='doublecircle', style='filled', fillcolor='lightblue')
            else:
                self.dot.node(str(state), shape='circle', style='filled', fillcolor='lightgray')
        
        # 添加转移边 - 直接绘制，不合并
        print("正在添加NFA边:")
        for (from_state, symbol), to_states in transitions.items():
            for to_state in to_states:
                print(f"  添加边: {from_state} --{symbol}--> {to_state}")
                if symbol == 'ε':
                    self.dot.edge(str(from_state), str(to_state), label=symbol, 
                                style='dashed', color='blue', fontcolor='blue', fontsize='12')
                else:
                    self.dot.edge(str(from_state), str(to_state), label=symbol, fontsize='12')
    
    def _render(self, output_file: str):
        """渲染图像"""
        os.makedirs('output', exist_ok=True)
        output_path = f'output/{output_file}'
        try:
            self.dot.render(output_path, format='png', cleanup=True, engine='dot')
            print(f"✅ NFA可视化图已保存到: {output_path}.png")
            return True
        except Exception as e:
            print(f"❌ 渲染失败: {e}")
            return False

def create_nfa_with_data(nfa_data):
    """根据提供的数据创建NFA可视化"""
    try:
        print("\n开始解析NFA数据...")
        
        # 解析输入数据
        states = NFAVisualizer.parse_states(nfa_data['states'])
        alphabet = NFAVisualizer.parse_alphabet(nfa_data['alphabet'])
        
        # 优先使用直接提供的字典格式
        if 'transitions_dict' in nfa_data:
            print("使用直接提供的NFA转换字典")
            transitions = nfa_data['transitions_dict']
        else:
            print("❌ NFA数据中缺少转换字典")
            return False
        
        start_state = NFAVisualizer.parse_start_state(nfa_data['start_state'])
        accept_states = NFAVisualizer.parse_accept_states(nfa_data['accept_states'])
        
        # 显示解析结果
        print("NFA解析结果:")
        print(f"  状态集合 Q: {states}")
        print(f"  字母表 Σ: {alphabet}")
        print(f"  转换函数 δ: {transitions}")
        print(f"  起始状态 q0: {start_state}")
        print(f"  接受状态 F: {accept_states}")
        
        # 验证转换表完整性
        print("验证转换表:")
        total_edges = 0
        for (from_state, symbol), to_states in transitions.items():
            print(f"  状态 {from_state} --{symbol}--> {to_states}")
            total_edges += len(to_states)
        print(f"  总共 {total_edges} 条边")
        
        # 创建可视化
        print("开始创建NFA可视化...")
        visualizer = NFAVisualizer()
        success = visualizer.visualize_from_data(
            states, alphabet, transitions, start_state, accept_states, nfa_data['output_file']
        )
        
        return success
    except Exception as e:
        print(f"❌ 解析NFA数据时出错: {e}")
        return False